/**
 * 
 */
package de.uni_postdam.ling.tcl.smc_lda;

import java.io.Reader;
import java.util.ArrayList;
import java.util.Date;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import org.apache.commons.math3.random.RandomGenerator;
import org.apache.commons.math3.random.Well19937a;
import de.uni_postdam.ling.tcl.util.ArraySampler;

/**
 * @author christoph
 *
 */
public class SequentialSampler
{
	/**
	 * 
	 */
	private ArrayList<Particle> active = new ArrayList<>();
	
	/**
	 * 
	 */
	private ArrayList<Particle> passive = new ArrayList<>();
	
	/**
	 * 
	 */
	private final Semaphore coordinator = new Semaphore(0);
	
	/**
	 * 
	 */
	private final ArrayList<Resampler> resamps = new ArrayList<>();
	
	/**
	 * 
	 */
	private final ArrayList<Sampler> samps = new ArrayList<>();

	/**
	 * 
	 */
	private ExecutorService es;
	
	/**
	 * 
	 * @return
	 */
	public SamplerGenerator makeSamplerGenerator()
	{
		return new SamplerGenerator();
	}
	
	/**
	 * 
	 * @param seed 
	 * @param historySize 
	 * @param gibbsRounds 
	 * @param topicWordSmoothing 
	 * @param numberOfTopics 
	 * @param topicSmoothing 
	 * @param rg
	 */
	SequentialSampler(int particles, long seed,
					int historySize, int gibbsRounds, double topicWordSmoothing,
					int numberOfTopics, double topicSmoothing)
	{
		RandomGenerator main = new Well19937a(seed);
		
		for(int i=0;i<particles;++i)
		{
			RandomGenerator rg = new Well19937a(main.nextLong());
			this.active.add(new Particle(numberOfTopics, topicSmoothing,
								topicWordSmoothing, gibbsRounds, historySize, rg));
			
			rg = new Well19937a(main.nextLong());
			this.passive.add(new Particle(numberOfTopics, topicSmoothing,
								topicWordSmoothing, gibbsRounds, historySize, rg));
			
			this.samps.add(new Sampler(i));
			
			rg = new Well19937a(main.nextLong());
			this.resamps.add(new Resampler(i, rg));
		}
	}
	
	/**
	 * 
	 * @param threads
	 * @param ds
	 * @throws InterruptedException 
	 */
	public void run(int threads, DocumentSource ds,
			int initializationDocs, int initializationSteps) throws InterruptedException
	{
		ds.ensureLookAheadSize(initializationDocs,this.active.get(0).getHistorySize());
		
		this.es = Executors.newFixedThreadPool(threads);
		
		initialize(ds, initializationDocs, initializationSteps);
		
		double[] weights = new double[this.active.size()];
		
		for(int i=0;i<this.resamps.size();++i)
		{
			this.resamps.get(i).setCwf(weights);
			this.samps.get(i).setDs(ds);
		}
		
		do
		{
			sample(ds, weights);
		}while(ds.hasMoreData());
		
		es.shutdown();
	}

	/**
	 * 
	 * @param ds
	 * @param weights
	 * @throws InterruptedException
	 */
	private void sample(DocumentSource ds, double[] weights)
			throws InterruptedException
	{
		ArrayList<Particle> inter = this.active;
		this.active = this.passive;
		this.passive = inter;
		
		for(int i=0;i<this.passive.size();++i)
		{
			weights[i] = this.passive.get(i).getLogImportanceWeight();
		}
		
		this.resamps.get(0).samp.turnIntoCWF(weights);
		
		for(int i=0;i<this.resamps.size();++i)
		{
			es.submit(this.resamps.get(i));
		}
		
		try
		{
			this.coordinator.acquire(this.active.size());
		}
		catch(InterruptedException ie)
		{
			es.shutdownNow();
			throw ie;
		}
		
		ds.removeBlock(this.active.get(0).getHistorySize());
	}

	/**
	 * 
	 * @param ds
	 * @param initializationDocs
	 * @param initializationSteps
	 * @throws InterruptedException
	 */
	private void initialize(DocumentSource ds, int initializationDocs,
			int initializationSteps) throws InterruptedException
	{
		for(int i=0;i<this.active.size();++i)
		{
			final int pos = i;
			es.submit(() -> 
					{
						this.active.get(pos).initialize(initializationDocs, initializationSteps, ds);
						this.coordinator.release();
					});
		}
		
		try
		{
			this.coordinator.acquire(this.active.size());
		}
		catch(InterruptedException ie)
		{
			es.shutdownNow();
			throw ie;
		}
		
		ds.removeBlock(initializationDocs);
	}
	
	/**
	 * 
	 * @author christoph
	 *
	 */
	private class Sampler implements Runnable
	{
		/**
		 * 
		 */
		private final int pos;
		
		/**
		 * 
		 */
		private DocumentSource ds;

		/**
		 * @param pos
		 */
		public Sampler(int pos)
		{
			super();
			this.pos = pos;
		}
		
		/**
		 * @param ds the ds to set
		 */
		private void setDs(DocumentSource ds)
		{
			this.ds = ds;
		}

		@Override
		public void run()
		{
			active.get(pos).expand(ds);
		}
	}
	
	
	private class Resampler implements Runnable
	{
		/**
		 * 
		 */
		private final int pos;
		
		/**
		 * 
		 */
		private double[] cwf;
		
		/**
		 * 
		 */
		private final ArraySampler samp;

		/**
		 * 
		 * @param rg
		 */
		private Resampler(int pos, RandomGenerator rg)
		{
			this.samp = new ArraySampler(rg);
			this.pos = pos;
		}
		
		/**
		 * @param cwf the cwf to set
		 */
		private void setCwf(double[] cwf)
		{
			this.cwf = cwf;
		}

		@Override
		public void run()
		{
			active.get(pos).rejuvenate(passive, cwf, samp);
			samps.get(pos).run();
			coordinator.release(1);
		}
	}
	
	/**
	 * 
	 * @author Christoph Teichmann
	 *
	 */
	public static class SamplerGenerator
	{
		
		private SamplerGenerator()
		{}
		
		/**
		 * 	
		 */
		private int	particles = 100;
		
		/**
		 * 
		 */
		private long	seed = new Date().getTime();
		
		/**
		 * 
		 */
		private int	historySize = 5;
		
		/**
		 * 
		 */
		private double	topicSmoothing = 0.1;
		
		/**
		 * 
		 */
		private int	numberOfTopics = 100;
		
		/**
		 * 
		 */
		private double	topicWordSmoothing = 10.0;
		
		/**
		 * 
		 */
		private int	gibbsRounds = 10;

		/**
		 * @param particles the particles to set
		 */
		public void setParticles(int particles)
		{
			this.particles = particles;
		}

		/**
		 * @param seed the seed to set
		 */
		public void setSeed(long seed)
		{
			this.seed = seed;
		}

		/**
		 * @param historySize the historySize to set
		 */
		public void setHistorySize(int historySize)
		{
			this.historySize = historySize;
		}

		/**
		 * @param topicSmoothing the topicSmoothing to set
		 */
		public void setTopicSmoothing(double topicSmoothing)
		{
			this.topicSmoothing = topicSmoothing;
		}

		/**
		 * @param numberOfTopics the numberOfTopics to set
		 */
		public void setNumberOfTopics(int numberOfTopics)
		{
			this.numberOfTopics = numberOfTopics;
		}

		/**
		 * @param topicWordSmoothing the topicWordSmoothing to set
		 */
		public void setTopicWordSmoothing(double topicWordSmoothing)
		{
			this.topicWordSmoothing = topicWordSmoothing;
		}

		/**
		 * @param gibbsRounds the gibbsRounds to set
		 */
		public void setGibbsRounds(int gibbsRounds)
		{
			this.gibbsRounds = gibbsRounds;
		}

		/**
		 * 
		 * @return
		 */
		public SequentialSampler make()
		{
			return new SequentialSampler(particles, seed, historySize,
							gibbsRounds, topicWordSmoothing, numberOfTopics, topicSmoothing);
		}
	}

	/**
	 * 
	 * @return
	 */
	public static SamplerGenerator generator()
	{
		return new SamplerGenerator();
	}

	/**
	 * 
	 * @param r1
	 */
	public void writeToFile(Reader r1)
	{
		// TODO Auto-generated method stub
		
	}
}
